from sqlalchemy import create_engine, Column, Integer, String
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker
from sqlalchemy.orm.session import Session
from models import Base 
from config import DATABASE_URL


engine = create_engine(DATABASE_URL)

class CustomSession(Session):
    def save(self, instance):
        """
        保存单个模型实例到数据库
        :param instance: SQLAlchemy 模型实例
        :return: 保存的实例
        """
        try:
            self.add(instance)
            self.commit()
            self.refresh(instance)
            return instance
        except Exception as e:
            self.rollback()
            raise e

    def save_all(self, instances):
        """
        保存多个模型实例到数据库
        :param instances: SQLAlchemy 模型实例列表
        :return: 保存的实例列表
        """
        try:
            self.add_all(instances)
            self.commit()
            for instance in instances:
                self.refresh(instance)
            return instances
        except Exception as e:
            self.rollback()
            raise e


# 使用自定义的 Session 类创建 SessionLocal
SessionLocal = sessionmaker(
    autocommit=False,
    autoflush=False,
    bind=engine,
    class_=CustomSession
)
def get_db():
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()
# 创建数据库表
Base.metadata.create_all(bind=engine)

